package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"database/sql"
	"encoding/pem"
	"fmt"
	"log"

	_ "github.com/go-sql-driver/mysql"
)

func GenerateRSAKeyPair(keySize int) (string, string, error) {
	// 生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		return "", "", err
	}

	// 序列化私钥为 PEM 格式
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privateKeyBytes,
	})

	// 获取公钥并序列化为 PEM 格式
	publicKey := privateKey.PublicKey
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		return "", "", err
	}
	publicKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: publicKeyBytes,
	})

	return string(privateKeyPEM), string(publicKeyPEM), nil
}

func StoreKeysInDB(privateKey, publicKey string) error {
	// 连接 MySQL 数据库
	dsn := "root:123456@tcp(127.0.0.1:3306)/go_admin?charset=utf8mb4&parseTime=True&loc=Local"
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return err
	}
	defer db.Close()

	// 插入密钥到数据库
	_, err = db.Exec("INSERT INTO sys_client (private_key, public_key, app_code) VALUES (?, ?,?)", privateKey, publicKey, "ceshi")
	if err != nil {
		return err
	}
	return nil
}

func main() {
	privateKeyPEM, publicKeyPEM, err := GenerateRSAKeyPair(2048)
	if err != nil {
		fmt.Println("Error generating key pair:", err)
		return
	}
	fmt.Println("Private Key:", privateKeyPEM)
	fmt.Println("Public Key:", publicKeyPEM)
	err = StoreKeysInDB(privateKeyPEM, publicKeyPEM)
	if err != nil {
		log.Fatalf("Error storing keys in database: %v", err)
	}
	fmt.Println("Keys stored successfully!")
}

func GetKeysFromDB(db *sql.DB) (string, string, error) {
	var privateKey, publicKey string
	err := db.QueryRow("SELECT private_key, public_key FROM rsa_keys LIMIT 1").Scan(&privateKey, &publicKey)
	if err != nil {
		return "", "", err
	}
	return privateKey, publicKey, nil
}
